import os
import random
import time

import numpy as np
import webdataset as wds


def set_seed(seed: int, rank: int = 0):
    """
    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch` and/or `tf` (if installed).

    Args:
        seed (`int`): The seed to set.
    """
    from nextstep.utils.import_utils import is_torch_available

    random.seed(seed + rank)
    np.random.seed(seed + rank)
    if is_torch_available():
        import torch

        torch.manual_seed(seed + rank)
        # ^^ safe to call this function even if cuda is not available
        torch.cuda.manual_seed_all(seed + rank)
        torch.backends.cudnn.deterministic = True
    os.environ["PYTHONHASHSEED"] = str(seed + rank)


def make_seed(*args):
    """
    Generate a seed value based on the hashed values of input arguments.

    This function takes any number of arguments, hashes each argument, and combines them into a single seed value.
    The seed is generated by iteratively multiplying the current seed by 31, adding the hash of the next argument,
    and taking the result modulo 2^31. This ensures the seed stays within the range of a 32-bit integer.

    Args:
        *args: Variable length argument list to be hashed and combined.

    Returns:
        int: The final seed value.
    """
    seed = 0
    for arg in args:
        seed = (seed * 31 + hash(arg)) & 0x7FFFFFFF
    return seed

